# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
from absl import app, flags

from coin.experimental.xguo.fruit.grape.price_graph import PriceGraph, OrderSide
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.experimental.xguo.fruit.base.strategy_base import (
    BaseStrategy,
    BaseStrategyConfig,
)
from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.util.util import FillStats
from coin.experimental.xguo.fruit.grape.trimm_sim import TriMMSimStrategy

FLAGS = flags.FLAGS


class TriMMStrategy(BaseStrategy):
  def __init__(self, config):
    super().__init__(config)
    self._pricer = PriceGraph(self.product_list, ignore_fee=True)
    self._fill_stats = collections.defaultdict(lambda: None)
    # self._sim_strat = TriMMSimStrategy(config)

    trade_products = []
    for product in self.product_list:
      if product.quote.currency == 'BTC':
        trade_products.append(product)
    assert trade_products

    self._logger.info('trade products: %s', trade_products)
    self._trade_products = trade_products

  def is_trade_product(self, product):
    return product in self._trade_products

  def prepare_fill_stat(self, product):
    if self._dry_run:
      return

    fee_bp = get_fee_rate(product)[0]
    self._fill_stats[product] = FillStats(product, fee_bp=fee_bp)
    self._scheduler.add_periodic_task(
        period=30,
        job_func=self._fill_stats[product].print,
        msg='Real PNL from fill')

  def prepare_pricers(self):
    return
    for products in self.product_triad_list:
      trade_symbol = pricer_info['trade_symbol']
      ref_symbol = pricer_info['ref_symbol']
      executor = self.get_executor(trade_symbol)
      exe_config = self.strategy_config.executor_config[trade_symbol]['proto_config']
      config = TmmPricerConfig()
      config.trade_symbol = trade_symbol
      config.ref_symbol = ref_symbol
      config.get_pos = lambda: executor.og.get_balance(trade_symbol.base)
      config.max_pos = exe_config.max_pos
      config.min_pos = exe_config.min_pos
      config.force_fire_cap_qty = exe_config.lot_size
      self._pricers[base]['pricer'] = SmmPricer(config)
      self._logger.info('Add smm pricer for %s', base)

  def install_periodic_task(self):
    self._scheduler.add_periodic_task(
        period=900,
        job_func=self.write_log,
    )
    self._scheduler.add_periodic_task(
        period=30,
        job_func=self.dump_pnl_balance,
    )
    self._scheduler.start()

  def run_before_start(self):
    self.install_periodic_task()
    self.register_balance_printer()

    self.prepare_pricers()
    for prod in self.product_list:
      self.prepare_fill_stat(prod)

    for pricer_info in self._pricers.values():
      self._logger.info('update tick for korean exchanges!')
      pricer_info['pricer'].update_tick()

    for og in self.order_gateway_list:
      self._logger.info('Update tick formatter for korean exchanges!')
      og.renew_order_formatter()

    for executor in self.executor_list:
      self._logger.info('Update tick for executor!')
      executor.update_tick()

  def write_log(self):
    if not self._pricer.is_ready():
      return
    self._strategy_logger.write_request()

    for og in self.order_gateway_list:
      og.og_logger.write_account_request()

    reserve_proto = self.strategy_config.strategy_core_config['reserve_proto']
    for proto in reserve_proto.values():
      self._strategy_logger.write_reserve(proto)

    for config in self.strategy_config.executor_config.values():
      self._strategy_logger.write_executor_config(config['proto_config'])

  def register_balance_printer(self):
    if self._dry_run:
      return

    assert len(self.order_gateway_list) == 1
    og = self.order_gateway_list[0]
    subreq = self.strategy_config.strategy_core_config['subreq_list'][0]
    printer_cls = get_pnl_balance_printer_cls(subreq)
    self._balance_printer = printer_cls(
        self._trade_products,
        self._book_map,
        reserve_map,
    )
    self._balance_printer.register_order_gateway(og)

  def dump_pnl_balance(self):
    if not self.is_initialized():
      return

    if not self._pricer.is_ready():
      return

    if self._balance_printer is not None:
      self._balance_printer.print_pnl_balance()
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)

  def on_book(self, product, book):
    # self._sim_strat.on_common_book(book)
    if not self.is_initialized():
      return
    self._pricer.update(product, book)
    for x in self._pricer.paths:
      print(x)
    return
    if not self._pricer.is_ready():
      return

    fill_stat = self._fill_stats.get(product)
    if fill_stat:
      fill_stat.update_book(book)

    pricer = self.get_pricer(product)
    pricer.update(product, book)
    if not pricer.is_ready():
      return

    orders = pricer.gen_order()
    self._writer.write(orders, book.timestamp)
    for prod, order in orders.items():
      executor = self.get_executor(prod)
      agg_buy_order = orders.get('agg_buy')
      agg_sell_order = orders.get('agg_sell')
      mm_order = orders.get('mm')
      self._writer.write(orders, book.timestamp)

      if self._dry_run:
        if self._dry_run == 'VERBOSE':
          self.log_json_message(orders, msg="Submit orders: ")
        return

      if agg_buy_order:
        agg_qty = agg_buy_order['qty']
        if agg_qty > 1e-4:
          executor.force_fire_buy_order(agg_buy_order['price'], agg_qty)

      if agg_sell_order:
        agg_qty = agg_sell_order['qty']
        if agg_qty > 1e-4:
          executor.force_fire_sell_order(agg_sell_order['price'], agg_qty)

      if mm_order:
        executor.manage_mm_orders(**mm.params)
        executor.manage_mm_orders(
              buy_price=mm_order['pass_buy_price'],
              sell_price=mm_order['pass_sell_price'],
              buy_price_pull=mm_order['buy_price_pull'],
              sell_price_pull=mm_order['sell_price_pull'],
              buy_price_push=mm_order['buy_price_push'],
              sell_price_push=mm_order['sell_price_push'])
    self._writer.write(convert_book_to_dict(book, product=product), book.timestamp)

  def on_trade(self, product, trade_event):
    self._sim_strat.on_trade(product, trade_event)

  def on_fill(self, fill_price, fill_qty, order):
    self._logger.info('ORDER_FILL: %s, %s, %s', order.order_id, fill_price, fill_qty)
    self._fill_stats[order.product].update(fill_price, fill_qty, order)
    fill_dict = {
        'type': 'fill',
        'fill_price': fill_price,
        'fill_qty': fill_qty,
        'order': order._asdict(),
    }
    self._writer.write(fill_dict, get_timestamp())


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = TriMMStrategy(config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')
  app.run(main)
